import os
import sys
import cv2
import dlib
import shutil
import corecode.api as corecode
import numpy as np
np.set_printoptions(threshold=32, edgeitems=10)
try:
   import cPickle as pickle
except:
   import pickle

def folderContents(dire):
    for dirpath, dirnames, filenames in os.walk(dire):
        return dirpath, dirnames, filenames

def loadPickle(ff):
    with open(ff, "rb") as openfile:
         pData = (pickle.load(openfile, encoding='latin1'))  
         return pData
    ff.close()


def getEmbeddings(ebl="embeddings"):
    mainNameList = []
    mainEmbeList = []

    cwd = os.getcwd()

    print ("Loading embeddings....")
    mainLoc, fConts, fFiles = folderContents(cwd + "/" + ebl + "/")
    for i in fConts: 
        mainLoc_2, fConts_2, fFiles_2 = folderContents(mainLoc + "/" + i)
        fFiles_2b = filter(lambda x : x != '*.jpg', fFiles_2)
        for i in fFiles_2b:
            if not any(value in (i.lower()) for value in (".jpeg", ".jpg", ".gif", ".png")):
               totalString = (mainLoc_2 + "/" + i)

               name = totalString.split("//")[-1].split("/")[0]
               nameFL = name.replace(".", " ").title()
               eT = totalString.split("//")[-1].split(".")[1].split("/")[1]
               nameMain = (nameFL + " - EmbeddingID[{}]".format(eT))
               mainNameList.append(nameMain)

               ebData = loadPickle(totalString)
               mainEmbeList.append(ebData)

               print  (nameMain + str(ebData) + '\n')
    print ("Done loading embeddings.")
    return mainNameList, mainEmbeList

def identFacesInPictures(names, known_faces, main_image):
    doneList = []               
    face_locations = corecode.face_locations(main_image)
    face_encodings = corecode.face_encodings(main_image, face_locations)    
    for i, ((top, right, bottom, left), atv_encoding) in enumerate(zip(face_locations, face_encodings)):
        for ii, (name, known_face) in enumerate(zip(names, known_faces)): 
           if (str(top) + str(left)) in doneList: break;

           match = corecode.compare_faces([known_face], atv_encoding, 0.55)
              
           if match[0] == True:
                print ("{} - {} - {}".format(i,ii,name))
                cv2.rectangle(main_image, (left, top), (right, bottom), (255, 0, 0), 2) 
                # Draw a label with a name below the face
                cv2.rectangle(main_image, (left, bottom - 35), (right, bottom), (255, 0, 0), cv2.FILLED)
                font = cv2.FONT_HERSHEY_DUPLEX
                cv2.putText(main_image, name, (left + 6, bottom - 6), font, 1.0, (0, 0, 0), 1)

                doneList.append(str(top) + str(left))

           elif match[0] == False:
                name = "Other"
                cv2.rectangle(main_image, (left, top), (right, bottom), (0, 0, 255), 2) 
                # Draw a label with a name below the face
                cv2.rectangle(main_image, (left, bottom - 35), (right, bottom), (0, 0, 255), cv2.FILLED)
                font = cv2.FONT_HERSHEY_DUPLEX
                cv2.putText(main_image, name, (left + 6, bottom - 6), font, 1.0, (255, 255, 255), 1)                    
    return main_image


def main():
    vF = sys.argv[1]
    if vF is None:
            print ('Not a video')
            sys.exit(1)
    if ".mp4" not in str(vF):
            print ('Not an mp4.')
            sys.exit(1)
    video_capture = cv2.VideoCapture(vF)
    _names, _embe = getEmbeddings()

    while True:
        # Grab a single frame of video
        ret, frame = video_capture.read()
        print(str(type(frame)))

        fFrame = identFacesInPictures(_names, _embe, frame)
        
 
        # Display the resulting image
        cv2.imshow('Video', fFrame)

        # Hit 'q' on the keyboard to quit!
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # Release handle to the webcam
    video_capture.release()
    cv2.destroyAllWindows()


main()
